#!/usr/bin/env python3

import sys
import os
import time
from p4runtime_sh.shell import *
from p4runtime_sh import p4runtime_lib

def program_switch(p4info_file_path, bmv2_file_path, switch_address):
    # 创建 P4Runtime 连接
    client = p4runtime_lib.bmv2.Bmv2SwitchConnection(
        name='s1',
        address=switch_address,
        device_id=0,
        proto_dump_file='logs/p4runtime_requests.txt'
    )

    # 设置管道配置
    client.SetForwardingPipelineConfig(
        p4info=p4runtime_lib.helper.P4Info(p4info_file_path),
        bmv2_json_file_path=bmv2_file_path
    )

    print("Installing ACL rules using P4Runtime...")

    try:
        # 规则1: 丢弃所有目标端口为80的UDP数据包
        table_entry = p4runtime_lib.helper.TableEntry('MyIngress.acl')(
            action=p4runtime_lib.helper.Action('MyIngress.drop')
        )
        table_entry.match['hdr.udp.dstPort'] = p4runtime_lib.helper.Ternary('80', '0xFFFF')
        table_entry.priority = 10
        client.WriteTableEntry(table_entry)
        print("Installed rule: drop UDP packets with dstPort=80")

        # 规则2: 丢弃所有目标IP为10.0.1.4的数据包
        table_entry = p4runtime_lib.helper.TableEntry('MyIngress.acl')(
            action=p4runtime_lib.helper.Action('MyIngress.drop')
        )
        table_entry.match['hdr.ipv4.dstAddr'] = p4runtime_lib.helper.Ternary('10.0.1.4', '0xFFFFFFFF')
        table_entry.priority = 10
        client.WriteTableEntry(table_entry)
        print("Installed rule: drop packets with dstIP=10.0.1.4")

        # 安装转发规则
        print("Installing forwarding rules...")

        # h1: 10.0.1.1 -> port 1, MAC: 00:00:00:00:01:01
        table_entry = p4runtime_lib.helper.TableEntry('MyIngress.ipv4_lpm')(
            action=p4runtime_lib.helper.Action('MyIngress.ipv4_forward',
                                              dstAddr='00:00:00:00:01:01',
                                              port=1)
        )
        table_entry.match['hdr.ipv4.dstAddr'] = p4runtime_lib.helper.Lpm('10.0.1.1', 32)
        client.WriteTableEntry(table_entry)
        print("Installed forwarding rule for h1 (10.0.1.1)")

        # h2: 10.0.1.2 -> port 2, MAC: 00:00:00:00:01:02
        table_entry = p4runtime_lib.helper.TableEntry('MyIngress.ipv4_lpm')(
            action=p4runtime_lib.helper.Action('MyIngress.ipv4_forward',
                                              dstAddr='00:00:00:00:01:02',
                                              port=2)
        )
        table_entry.match['hdr.ipv4.dstAddr'] = p4runtime_lib.helper.Lpm('10.0.1.2', 32)
        client.WriteTableEntry(table_entry)
        print("Installed forwarding rule for h2 (10.0.1.2)")

        # h3: 10.0.1.3 -> port 3, MAC: 00:00:00:00:01:03
        table_entry = p4runtime_lib.helper.TableEntry('MyIngress.ipv4_lpm')(
            action=p4runtime_lib.helper.Action('MyIngress.ipv4_forward',
                                              dstAddr='00:00:00:00:01:03',
                                              port=3)
        )
        table_entry.match['hdr.ipv4.dstAddr'] = p4runtime_lib.helper.Lpm('10.0.1.3', 32)
        client.WriteTableEntry(table_entry)
        print("Installed forwarding rule for h3 (10.0.1.3)")

        # h4: 10.0.1.4 -> port 4, MAC: 00:00:00:00:01:04
        table_entry = p4runtime_lib.helper.TableEntry('MyIngress.ipv4_lpm')(
            action=p4runtime_lib.helper.Action('MyIngress.ipv4_forward',
                                              dstAddr='00:00:00:00:01:04',
                                              port=4)
        )
        table_entry.match['hdr.ipv4.dstAddr'] = p4runtime_lib.helper.Lpm('10.0.1.4', 32)
        client.WriteTableEntry(table_entry)
        print("Installed forwarding rule for h4 (10.0.1.4)")

        print("Successfully programmed switch with P4Runtime!")

    except Exception as e:
        print(f"Error programming switch: {e}")
        return False

    return True

if __name__ == '__main__':
    if len(sys.argv) != 4:
        print("Usage: python3 program_switch.py <p4info_file> <bmv2_json_file> <switch_address>")
        sys.exit(1)

    p4info_file = sys.argv[1]
    bmv2_json_file = sys.argv[2]
    switch_addr = sys.argv[3]

    if not os.path.exists(p4info_file):
        print(f"P4Info file not found: {p4info_file}")
        sys.exit(1)

    if not os.path.exists(bmv2_json_file):
        print(f"BMv2 JSON file not found: {bmv2_json_file}")
        sys.exit(1)

    success = program_switch(p4info_file, bmv2_json_file, switch_addr)
    sys.exit(0 if success else 1)
